#!/usr/bin/env python3
# -*- python -*-

#
# Copyright (C)  2020 ifm electronic, gmbh
# Copyright (C)  2017 Love Park Robotics, LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distribted on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

#
# Usage: ifm3d-dpkg-deps.py <foo>.deb
#
# Looks inside <foo>.deb. Finds all *.so files and computes the set of debian
# dependencies for which the shared objects of the passed in deb depend. It
# repackages <foo>.deb with these dependencies spelled out in the debian control
# file.
#

from __future__ import print_function
import argparse
import fileinput
import fnmatch
import os
import re
import shutil
import subprocess
import sys
import tempfile

VERSION = "@IFM3D_VERSION@"

class DpkgDeps(object):
    """Auto-computes the dependencies for a debian package and repackages it
    with the newly computed dependencies integral to the control file. It
    exploits some a-priori information about the ifm3d development
    environment.
    """

    def __init__(self, deb_file):
        self.deb_file_in_ = deb_file
        self.deb_file_tmp_ = None
        self.deb_file_control_ = None
        self.so_deps_ = []
        self.deb_deps_ = {} # deb => version
        self.deb_file_mod_ = "%s.mod.deb" % self.deb_file_in_

    def run(self):
        """Runs the debian dependency computation pipeline"""
        print("Processing input file: %s" % self.deb_file_in_)

        self.deb_file_tmp_ = tempfile.mkdtemp(prefix='ifm3d-deb-')
        self.deb_file_control_ = "%s/DEBIAN" % self.deb_file_tmp_
        print("Working directory is: %s" % self.deb_file_tmp_)

        retval = self.extract_deb()
        if retval != 0:
            return retval

        self.so_deps_ = self.generate_so_deps()
        self.deb_deps_ = self.generate_deb_deps()
        if len(self.deb_deps_) > 0:
            self.update_control_file()
        self.repackage_deb()

        print("Removing: %s" % self.deb_file_in_)
        os.remove(self.deb_file_in_)

        print("Renaming: %s --> %s" % (self.deb_file_mod_, self.deb_file_in_))
        os.rename(self.deb_file_mod_, self.deb_file_in_)

        print("Removing tmp directory: %s" % self.deb_file_tmp_)
        shutil.rmtree(self.deb_file_tmp_)

        return 0

    def repackage_deb(self):
        """Repackages the debian file with the new control information."""
        print("Building new deb: %s" % self.deb_file_mod_)
        out = subprocess.check_output(["dpkg", "-b",
                                       self.deb_file_tmp_,
                                       self.deb_file_mod_])

    def update_control_file(self):
        """Updates the debian control file"""
        print("Updating control file with dependencies...")
        with open("%s/control" % self.deb_file_control_, "a") as file:
            file.write("Depends: ")
            n = len(self.deb_deps_.keys())
            i = 0
            for deb, version in self.deb_deps_.items():
                if "ifm3d" in deb:
                    file.write("%s (= %s)" % (deb, version))
                else:
                    file.write("%s (>= %s)" % (deb, version))
                if i < n - 1:
                    file.write(", ")
                i += 1

        # now kill any blank lines in the file
        for line in fileinput.FileInput("%s/control" % self.deb_file_control_,
                                        inplace=1):
            if line.rstrip():
                print(line, end="")

        # Add final new line (dpkg requires this)
        with open("%s/control" % self.deb_file_control_, "a") as file:
            file.write("\n")

    def generate_deb_deps(self):
        """Based on the .so files in `self.so_deps_`, this creates a dict in
        `self.deb_deps_` that maps the needed debian packages to the associated
        version so that the .so dependencies are satisfied.
        """
        print("Computing debian dependencies for: %s" % self.deb_file_in_)
        debs = {}
        for so in self.so_deps_:
            if "ifm3d" in so:
                if "camera" in so:
                    debs["ifm3d-camera"] = VERSION
                elif "framegrabber" in so:
                    debs["ifm3d-framegrabber"] = VERSION
                elif "swupdater" in so:
                    debs["ifm3d-swupdater"] = VERSION
                elif "image" in so:
                    debs["ifm3d-image"] = VERSION
                elif "tools" in so:
                    debs["ifm3d-tools"] = VERSION
                elif "opencv" in so:
                    debs["ifm3d-opencv"] = VERSION
            else:
                try:
                    # Starting sometime between 18.04 and 20.04, /lib is a
                    # symlink to /usr/lib so must check both places
                    out = subprocess.check_output(["dpkg", "-S", os.path.realpath(so)]).decode('utf-8')
                except subprocess.CalledProcessError:
                    out = subprocess.check_output(["dpkg", "-S", so]).decode('utf-8')
                for line in out.split("\n"):
                    parts = line.split(":")
                    deb = parts[0]
                    if deb != "":
                        debs[deb] = 1

        deb_names = debs.keys()
        for deb in deb_names:
            if "ifm3d" in deb:
                continue

            out = subprocess.check_output(["dpkg", "-s", deb]).decode('utf-8')
            for line in out.split("\n"):
                if line.startswith("Version"):
                    m = re.search(r"Version:(.*)", line)
                    v = m.group(1).strip()
                    #if "~" in v:
                    #    v = v.split("~")[0]
                    debs[deb] = v

        return debs

    def generate_so_deps(self):
        """Returns a list of all the .so files this particular deb file relies
        upon.
        """
        print("Listing .so deps of: %s" % self.deb_file_in_)
        sos = []
        for root, dirnames, filenames in os.walk(self.deb_file_tmp_):
            for filename in fnmatch.filter(filenames, "*.so"):
                sos.append(os.path.join(root, filename))
        print("Found the following .so files: %s" % str(sos))

        so_deps = {}
        # This gives us the base name of the .so
        for so in sos:
            out = subprocess.check_output(["readelf", "-d", so]).decode('utf-8')
            for line in out.split("\n"):
                if "NEEDED" in line:
                    m = re.search(r"\[(.*)\]", line)
                    if m is not None:
                        so_deps[m.group(1)] = 1

        print(".so deps *prior* to full-path resolution:")
        print(so_deps)
        # now, resolve the full path to the .so
        for so in sos:
            out = subprocess.check_output(["ldd", so]).decode('utf-8')
            for line in out.split("\n"):
                parts = line.split()
                if len(parts) <= 1:
                    continue
                elif len(parts) < 3:
                    so_deps.pop(parts[0], None)
                else:
                    if ((parts[0] in so_deps) and
                            ("ifm3d" not in parts[0])):
                        so_deps[parts[2]] = so_deps.pop(parts[0])

        print(".so deps *after* to full-path resolution:")
        print(so_deps)

        retval = so_deps.keys()
        print("Found the following .so deps: %s" % retval)
        return retval

    def extract_deb(self):
        """Extracts the input deb file to the temp directory"""
        print("Extracting %s to %s" % (self.deb_file_in_, self.deb_file_tmp_))
        subprocess.check_call(["dpkg-deb", "-x", self.deb_file_in_,
                               self.deb_file_tmp_])

        print("Extracing control info to: %s" % self.deb_file_control_)
        subprocess.check_call(["dpkg-deb", "--control", self.deb_file_in_,
                               self.deb_file_control_])
        return 0

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("deb_path", type=str, nargs='*',
                        help="Path to deb or glob expression to multiple debs")
    args = parser.parse_args()

    retval = 0
    for deb in args.deb_path:
        dpkg = DpkgDeps(deb)
        retval = dpkg.run()
        if retval != 0:
            sys.stderr.write("Failed to create dependencies for: %s\n" % deb)
            break

    return retval

if __name__ == '__main__':
    sys.exit(main())
